from flask import Flask, Blueprint, request, jsonify, current_app, send_from_directory
import os
import uuid
import time
from pathlib import Path
import logging
from ultralytics import YOLO

# 创建Flask应用
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024  # 50MB

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建视频上传蓝图
video_bp = Blueprint('video', __name__)

# 允许的文件扩展名
ALLOWED_EXTENSIONS = {'mp4', 'webm', 'jpg', 'jpeg', 'png'}

# 全局模型变量
a1 = None
model_loaded = False


def load_model():
    """加载YOLO模型"""
    global a1, model_loaded

    try:
        # 方法1: 使用相对路径
        current_dir = Path(__file__).parent
        model_path = current_dir / 'best.pt'

        # 方法2: 如果上述路径不存在，尝试其他路径
        if not model_path.exists():
            model_path = current_dir.parent / 'models' / 'best.pt'

        # 方法3: 如果还是不存在，使用绝对路径但确保路径正确
        if not model_path.exists():
            # 请根据您的实际路径修改
            model_path = r'C:\Users\20798\xwechat_files\wxid_vn8zizbnoofj22_bdf5\msg\file\2025-11\src (2)\src\src\backend\routes\best.pt'

        logger.info(f"尝试加载模型: {model_path}")
        logger.info(f"模型文件存在: {os.path.exists(model_path)}")

        if os.path.exists(model_path):
            a1 = YOLO(str(model_path))
            model_loaded = True
            logger.info("YOLO模型加载成功!")
            logger.info(f"模型类别: {a1.names}")
        else:
            logger.error(f"模型文件不存在: {model_path}")
            a1 = None
            model_loaded = False

    except Exception as e:
        logger.error(f"YOLO模型加载失败: {e}")
        import traceback
        traceback.print_exc()
        a1 = None
        model_loaded = False


# 在模块导入时立即加载模型
logger.info("开始加载YOLO模型...")
load_model()


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@video_bp.route('/api/upload-file', methods=['POST'])
def upload_file():
    global a1, model_loaded

    # 每次请求都检查模型状态
    if a1 is None or not model_loaded:
        logger.warning("模型未加载，尝试重新加载...")
        load_model()

    try:
        # 检查模型是否加载
        if a1 is None or not model_loaded:
            logger.error("模型仍未加载，无法进行分析")

            if 'file' not in request.files:
                return jsonify({
                    'code': 400,
                    'message': '没有文件部分'
                }), 400

            file = request.files['file']

            if file.filename == '':
                return jsonify({
                    'code': 400,
                    'message': '未选择文件'
                }), 400

            if not allowed_file(file.filename):
                return jsonify({
                    'code': 400,
                    'message': '不支持的文件类型'
                }), 400

            # 即使模型未加载也保存文件
            upload_folder = current_app.config.get('UPLOAD_FOLDER', 'uploads')
            os.makedirs(upload_folder, exist_ok=True)

            file_extension = file.filename.rsplit('.', 1)[1].lower()
            timestamp = int(time.time())
            unique_filename = f"{timestamp}_{uuid.uuid4().hex[:8]}.{file_extension}"
            file_path = os.path.join(upload_folder, unique_filename)

            file.save(file_path)
            logger.info(f"文件保存成功: {file_path}")

            return jsonify({
                'code': 200,
                'message': '文件上传成功（模型未加载，无法分析）',
                'data': {
                    'filename': unique_filename,
                    'original_filename': file.filename,
                    'file_path': file_path
                }
            }), 200

        # 正常处理文件上传和预测
        if 'file' not in request.files:
            return jsonify({
                'code': 400,
                'message': '没有文件部分'
            }), 400

        file = request.files['file']

        if file.filename == '':
            return jsonify({
                'code': 400,
                'message': '未选择文件'
            }), 400

        if not allowed_file(file.filename):
            return jsonify({
                'code': 400,
                'message': '不支持的文件类型'
            }), 400

        upload_folder = current_app.config.get('UPLOAD_FOLDER', 'uploads')
        os.makedirs(upload_folder, exist_ok=True)

        file_extension = file.filename.rsplit('.', 1)[1].lower()
        timestamp = int(time.time())
        unique_filename = f"{timestamp}_{uuid.uuid4().hex[:8]}.{file_extension}"
        file_path = os.path.join(upload_folder, unique_filename)
        file.save(file_path)
        logger.info(f"文件保存成功: {file_path}")

        # 使用 YOLO 模型进行预测
        try:
            logger.info("开始模型预测...")
            results = a1(file_path, show=False, save=True)

            # 提取带框图像路径（核心修改：确保取到正确的结果路径）
            result = results[0]  # 单张图片预测结果通常在results[0]
            save_dir = result.save_dir  # 模型保存带框图像的目录
            input_filename = os.path.basename(file_path)  # 原始文件名（带框图像与原始文件名相同）
            detection_image_path = os.path.join(save_dir, input_filename)  # 带框图像完整路径

            logger.info(f"预测完成，带框图像保存路径: {detection_image_path}")

            predictions = []
            # 遍历结果提取预测信息（复用原有逻辑）
            for i, result in enumerate(results):
                logger.info(f"处理第 {i} 个结果")
                boxes = result.boxes
                if boxes is not None:
                    logger.info(f"检测到 {len(boxes)} 个目标")
                    for j, box in enumerate(boxes):
                        cls = int(box.cls[0])
                        conf = float(box.conf[0])
                        class_name = result.names[cls] if hasattr(result, 'names') else str(cls)
                        predictions.append({
                            'class': cls,
                            'confidence': conf,
                            'class_name': class_name
                        })
                else:
                    logger.info("未检测到目标")

            logger.info(f"最终预测结果: {len(predictions)} 个目标")

            # 返回数据：file_path替换为带框图像路径，其他字段格式不变
            return jsonify({
                'code': 200,
                'message': '文件上传成功并完成分析',
                'data': {
                    'filename': unique_filename,  # 保持原始生成的唯一文件名（便于前端关联）
                    'original_filename': file.filename,
                    'file_path': detection_image_path,  # 核心：返回带框图像路径
                    'predictions': predictions
                }
            }), 200

        except Exception as e:
            logger.error(f"YOLO预测错误: {str(e)}")
            import traceback
            traceback.print_exc()

            return jsonify({
                'code': 500,
                'message': f'分析失败: {str(e)}',
                'data': {
                    'filename': unique_filename,
                    'original_filename': file.filename,
                    'file_path': file_path  # 失败时仍返回原始文件路径
                }
            }), 500

    except Exception as e:
        logger.error(f"文件上传错误: {str(e)}")
        import traceback
        traceback.print_exc()

        return jsonify({
            'code': 500,
            'message': f'服务器错误: {str(e)}'
        }), 500


# 添加一个检查模型状态的接口
@video_bp.route('/api/model-status', methods=['GET'])
def model_status():
    return jsonify({
        'code': 200,
        'message': '模型状态查询成功',
        'data': {
            'model_loaded': model_loaded,
            'model_path': str(Path(__file__).parent / 'best.pt') if model_loaded else None
        }
    }), 200

# 得到检测后图片的静态访问接口
@video_bp.route('/runs/detect/<path:filename>')
def serve_detection_image(filename):
    # 获取 runs/detect 目录的绝对路径（动态适配项目结构）
    runs_detect_dir = os.path.join(os.getcwd(), 'runs', 'detect')
    # 允许前端通过 /runs/detect/xxx/xxx.jpg 访问该目录下的文件
    return send_from_directory(runs_detect_dir, filename)

# 注册蓝图
app.register_blueprint(video_bp)

if __name__ == '__main__':
    # 确保上传目录存在
    os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

    # 禁用重载器，避免模型重复加载
    app.run(host='0.0.0.0', port=5000, debug=True, use_reloader=False)